#include "icmp.h"
#include "buildpacket.h"
#include "portInfo.h"
#include "scan.h"
#include "setFlags.h"
#include <iostream>
#include <unistd.h>
#include <sys/socket.h>
#include <memory>
#include <string.h>
#include <thread>
#include "checksum.h"
portScan::portScan(int startPort,int endPort,char* ip,char* mode):startPort(startPort),endPort(endPort),ip(ip),mode(mode){
    localIp = getLocalIP();
}
void portScan::receive(int port){
    struct sockaddr_in src_addr;
    socklen_t src_addr_len = sizeof(src_addr);

    char buffer[BUFFER_SIZE];
    struct ip *ip_hdr;
    struct tcphdr *tcp_hdr;

    int rawSocket1;
    // 创建原始套接字，接收 TCP 数据包
    rawSocket1 = socket(AF_INET, SOCK_RAW, IPPROTO_TCP);
    if (rawSocket1 == -1) {
        perror("socket");
    }

    // 设置筛选条件，例如只接收来自特定源地址的数据包
    struct in_addr target_ip;
    if (inet_aton(ip, &target_ip) == 0) {
        perror("inet_aton");
        close(rawSocket1);
        return;
    }

    int flags=10;
    while (flags) {
        // 接收数据包
        int data_size = recvfrom(rawSocket1, buffer, BUFFER_SIZE, 0, (struct sockaddr *)&src_addr, &src_addr_len);
        if (data_size == -1) {
            perror("recvfrom");
            close(rawSocket1);
            return;
        }

        ip_hdr = (struct ip *)buffer;
        tcp_hdr = (struct tcphdr *)(buffer + (ip_hdr->ip_hl * 4));
        // 筛选特定源地址的数据包
        if (ip_hdr->ip_src.s_addr == target_ip.s_addr && ntohs(tcp_hdr->source) == port) {
            if(tcp_hdr->syn && tcp_hdr->ack){
                PortStatus state = OPEN;
                PortInfo info = {port, state};
                std::lock_guard<std::mutex> lock(open_ports_mutex);
                open_ports.push_back(info);
                std::cout<<port<<" is open"<<std::endl;
                close(rawSocket1);
                return;
            }
            if(tcp_hdr->rst){
                if(!strcmp(mode,"-s")|| !strcmp(mode,"-f") || !strcmp(mode,"-n") || !strcmp(mode,"-x")){
                    PortStatus state = CLOSED;
                    PortInfo info = {port, state};
                    std::lock_guard<std::mutex> lock(closed_ports_mutex);
                    closed_ports.push_back(info);
                    std::cout<<port<<" is closed"<<std::endl; 
                    close(rawSocket1);   
                    return;
                }
                if(strcmp(mode,"-sa") == 0){
                    PortStatus state = OPEN;
                    PortInfo info = {port, state};
                    std::lock_guard<std::mutex> lock(open_ports_mutex);
                    open_ports.push_back(info);
                    std::cout<<port<<" is open"<<std::endl;
                    close(rawSocket1);    
                    return;
                }             
            }
        }else {
            flags--;
        }
    }
    if(!strcmp(mode,"-f") || !strcmp(mode,"-n") || !strcmp(mode,"-x")){
        PortStatus state = OPEN_FILTERED;
        PortInfo info = {port, state};
        std::lock_guard<std::mutex> lock(open_filtered_ports_mutex);
        open_filtered_ports.push_back(info);
        std::cout<<port<<" is open or filtered"<<std::endl; 
        close(rawSocket1);
    } 
    if(strcmp(mode,"-s") == 0){
        PortStatus state = FILTERED;
        PortInfo info = {port, state};
        std::lock_guard<std::mutex> lock(filtered_ports_mutex);
        filtered_ports.push_back(info);
        //std::cout<<port<<" is filtered"<<std::endl; 
        close(rawSocket1);
    }
    if(strcmp(mode,"-sa")){
        PortStatus state = CLOSED_FILTERED;
        PortInfo info = {port, state};
        std::lock_guard<std::mutex> lock(closed_filtered_ports_mutex);
        closed_filtered_ports.push_back(info);
        std::cout<<port<<" is closed or filtered "<<std::endl;
        close(rawSocket1); 
    }

}
void portScan::sendPacket(int port){
    struct sockaddr_in destAddr;

    // 初始化目标地址
    destAddr.sin_family = AF_INET;
    destAddr.sin_port = htons(port); // 目标端口号
    inet_pton(AF_INET, this->ip, &(destAddr.sin_addr)); // 目标IP地址
    struct sockaddr_in server_addr;

    char packet[BUFFER_SIZE];
    memset(packet, 0, sizeof(packet));

    dataPacket *myDataPacket = (dataPacket *)packet;
    struct iphdr *ipHeader = &myDataPacket->ip;
    struct tcphdr *tcpHeader = &myDataPacket->tcp;
    char *data = myDataPacket->data;

    // 填充IP头部
    ipHeader->ihl = 5;
    ipHeader->version = 4;
    ipHeader->tos = 0;
    ipHeader->tot_len = htons(sizeof(dataPacket));
    ipHeader->id = htons(12345);
    ipHeader->frag_off = 0;
    ipHeader->ttl = 64;
    ipHeader->protocol = IPPROTO_TCP;
    ipHeader->check = 0; // 校验和会在发送前计算
    ipHeader->saddr = inet_addr(this->localIp); // 源IP地址
    ipHeader->daddr = destAddr.sin_addr.s_addr;
    ipHeader->check = ip_checksum((unsigned short *)&ipHeader, sizeof(struct iphdr));
    
    // 填充TCP头部
    tcpHeader->source = htons(12345); // 源端口号
    tcpHeader->dest = destAddr.sin_port;
    tcpHeader->seq = htonl(0x12345678);
    tcpHeader->ack_seq = 0;
    tcpHeader->doff = 5;
    //tcpHeader->syn = 1;
    //tcpHeader->ack = 1;

    if(strcmp(mode,"-s") == 0){
        synPortScan(tcpHeader);
    }else if(strcmp(mode,"-f") == 0){
        finPortScan(tcpHeader);
    }else if(strcmp(mode,"-a") == 0){
        ackPortScan(tcpHeader);
    }else if(strcmp(mode,"-n") == 0){
        nullPortScan(tcpHeader);
    }else if(strcmp(mode,"-sa") == 0){
        syn_ackPortScan(tcpHeader);
    }else if(strcmp(mode,"-x") == 0){
        xmasPortScan(tcpHeader);
    }else{
        std::cout<<"扫描模式非法"<<std::endl; 
        return;
    }
    tcpHeader->window = htons(512);
    tcpHeader->check = 0; // 校验和会在发送前计算
    tcpHeader->urg_ptr = 0;

    // 填充数据部分
    // memset(data, 'A', 0);
    // 计算校验和
    tcp_checksum(ipHeader, tcpHeader, data, 0);

    // 初始化 socket
    int rawSocket = socket(AF_INET, SOCK_RAW, IPPROTO_RAW);
    if (rawSocket == -1) {
        perror("socket");
        exit(1);
    }

    int optval = 1;
    int set = setsockopt(rawSocket, IPPROTO_IP, IP_HDRINCL, &optval, sizeof(optval));
    if (set == -1) {
        perror("set");
        close(rawSocket);
        exit(1);
    }

    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(12345);
    server_addr.sin_addr.s_addr = INADDR_ANY;
    bind(rawSocket, (struct sockaddr *)&server_addr, sizeof(server_addr));            

    std::thread recv_thread(&portScan::receive, this, port);

    int sent = sendto(rawSocket, myDataPacket, ntohs(ipHeader->tot_len), 0, (struct sockaddr *)&destAddr, sizeof(destAddr));
    if (sent == -1) {
        perror("sendto");
        close(rawSocket);
        exit(1);
    }
    // 检查是否发送了预期的字节数
    if (sent != ntohs(ipHeader->tot_len)) {
        printf("Warning: Not all bytes were sent. Sent: %d, Expected: %d\n", sent, ntohs(ipHeader->tot_len));
    }

    recv_thread.join();
}

void portScan::run(){

    for(int i=startPort;i<=endPort;i++){
        sendPacket(i);
    }


}

void scanPorts(int startPort, int endPort, char* ip, char* mode) {
    portScan *scan = new portScan(startPort, endPort, ip, mode);
    scan->run();
}

void allotThread(int startPort, int endPort, char*portScanMode,int threadCount,char* ip)
{
    int totalPorts = endPort - startPort + 1;
    int portsPerThread = totalPorts / threadCount;
    int remainingPorts = totalPorts % threadCount;

    std::vector<std::thread> threads;
    if(totalPorts >=threadCount){
            // 创建并启动线程
        for (int i = 0; i < threadCount; ++i) {
            int currentStartPort = startPort + i * portsPerThread;
            int currentEndPort = (i == threadCount - 1) ? endPort : (currentStartPort + portsPerThread - 1);

            // 如果还有剩余的端口未分配，则给前几个线程多分一些
            if (remainingPorts > 0) {
                currentEndPort++;
                remainingPorts--;
            }
            threads.emplace_back(scanPorts, currentStartPort, currentEndPort, ip, portScanMode);

        }
    }else{
        for (int i = startPort; i <= endPort; ++i){
            int currentPort = i;
            threads.emplace_back(scanPorts, currentPort, currentPort, ip, portScanMode);
        }
    }

    // 等待所有线程完成
    for (auto& t : threads) {
        t.join();
    }
}